#ifndef _CODEGEN_SYMBOL_H
#define _CODEGEN_SYMBOL_H

#include "llvm/IR/Function.h"
#include "llvm/IR/Type.h"
#include "llvm/IR/Value.h"
#include <unordered_map>
#include <memory>
#include <map>
#include <string>
#include <vector>

class Symbol {
public:
    // ALIAS means type alias, PTYPE means ALIAS or ARRAY or RECORD
    enum class Type {
        UNKNOWN, VARIABLE, CONSTANT, FUNCTION, ALIAS, ARRAY, RECORD, PTYPE
    };
    Symbol() = default;
    Symbol(Type T) : T(T) {}
    virtual ~Symbol() = default;
    Type type() const noexcept { return T; }
private:
    Type T;
};

class SymbolVariable : public Symbol {
public:
    SymbolVariable(llvm::Value* pointer, std::string type, std::shared_ptr<Symbol> ref)
        : Symbol(Symbol::Type::VARIABLE), pointer(pointer), type(type), ref(ref) {}
    llvm::Value* getValue() { return pointer; }
    std::string getTypeID() { return type; }
    std::shared_ptr<Symbol> getTypeSymbol() { return ref; }
private:
    llvm::Value* pointer;
    // Array or record
    // `type' for type identifier
    // `ref' for corresponding symbol
    std::string type;
    std::shared_ptr<Symbol> ref;
};

class SymbolConstant : public Symbol {
public:
    SymbolConstant(llvm::Value* value)
        : Symbol(Symbol::Type::CONSTANT), value(value) {}
    llvm::Value* getValue() { return value; }
private:
    llvm::Value* value;
};

class SymbolFunction : public Symbol {
public:
    SymbolFunction(llvm::Function* function, std::vector<bool> var)
        : Symbol(Symbol::Type::FUNCTION), function(function), var(var) {}
    llvm::Function* getFunction() { return function; }
    std::vector<bool> getVarPos() { return var; }
private:
    llvm::Function* function;
    std::vector<bool> var;
};

// Type alias
// Nested type alias not support
class SymbolAlias : public Symbol {
public:
    SymbolAlias(llvm::Type* type)
        : Symbol(Symbol::Type::ALIAS), type(type) {}
    llvm::Type* getType() { return type; }
private:
    // Only support simple/standard types
    llvm::Type* type;
};

// Only support one-dimensional array
class SymbolArray : public Symbol {
public:
    SymbolArray(llvm::ArrayType* array, int beg, int end)
        : Symbol(Symbol::Type::ARRAY), array(array), beg(beg), end(end) {}
    llvm::ArrayType* getArray() { return array; }
    int getIndexBegin() { return beg; }
    int getIndexEnd() { return end; }
private:
    llvm::ArrayType* array;
    int beg;
    int end;
};

class SymbolRecord : public Symbol {
public:
    SymbolRecord(llvm::StructType* record, std::vector<std::string> field)
        : Symbol(Symbol::Type::RECORD), record(record), field(field) {}
    llvm::StructType* getRecord() { return record; }
    std::vector<std::string> getField() { return field; }
private:
    llvm::StructType* record;
    // Record field type must be simple/standard types
    std::vector<std::string> field;
};

class SymbolTable {
public:
    SymbolTable() = default;
    std::shared_ptr<Symbol> find(std::string name) const {
        auto it = namedValues.find(name);
        if (it == namedValues.end()) {
            return nullptr;
        }
        return it->second;
    }
    int insert(std::string name, std::shared_ptr<Symbol> symbol) {
        if (namedValues.find(name) != namedValues.end()) {
            return -1;
        }
        namedValues.insert({ name, symbol });
        return 0;
    }
private:
    std::unordered_map<std::string, std::shared_ptr<Symbol>> namedValues;
};

class Scope {
public:
    Scope() = default;

    void push() { stack.push_back(std::make_shared<SymbolTable>()); }
    void pop() { stack.pop_back(); }

    // If depth()=1, variables and constants are global
    std::size_t depth() const noexcept { return stack.size(); }

    // Find symbol from symbol stack
    std::shared_ptr<Symbol> find(std::string name, Symbol::Type type = Symbol::Type::UNKNOWN) const {
        for (auto it = stack.rbegin(); it != stack.rend(); it++) {
            auto ret = (*it)->find(name);
            if (nullptr != ret) {
                auto rtype = ret->type();
                if (rtype == type
                    || Symbol::Type::UNKNOWN == type
                    || (Symbol::Type::PTYPE == type
                        && (Symbol::Type::ALIAS == rtype
                            || Symbol::Type::ARRAY == rtype
                            || Symbol::Type::RECORD == rtype))) {
                    return ret;
                }
                break;
            }
        }
        return nullptr;
    }

    // We do not allow any form of overloading
    // If is_local() returns false then it is okay to insert
    // Or you can use return value of insert() to detect
    bool is_local(std::string name) const {
        return nullptr == stack.back()->find(name) ? false : true;
    }

    int insertVariable(std::string name, llvm::Value* pointer,
        std::string type = "", std::shared_ptr<Symbol> ref = nullptr) {
        std::shared_ptr<SymbolVariable> symbol = std::make_shared<SymbolVariable>(pointer, type, ref);
        return stack.back()->insert(name, symbol);
    }
    int insertConstant(std::string name, llvm::Value* value) {
        std::shared_ptr<SymbolConstant> symbol = std::make_shared<SymbolConstant>(value);
        return stack.back()->insert(name, symbol);
    }
    int insertFunction(std::string name, llvm::Function* function, std::vector<bool> var) {
        std::shared_ptr<SymbolFunction> symbol = std::make_shared<SymbolFunction>(function, var);
        return stack.back()->insert(name, symbol);
    }
    int insertAlias(std::string name, llvm::Type* type) {
        std::shared_ptr<SymbolAlias> symbol = std::make_shared<SymbolAlias>(type);
        return stack.back()->insert(name, symbol);
    }
    int insertArray(std::string name, llvm::ArrayType* array, int beg, int end) {
        std::shared_ptr<SymbolArray> symbol = std::make_shared<SymbolArray>(array, beg, end);
        return stack.back()->insert(name, symbol);
    }
    int insertRecord(std::string name, llvm::StructType* record, std::vector<std::string> field) {
        std::shared_ptr<SymbolRecord> symbol = std::make_shared<SymbolRecord>(record, field);
        return stack.back()->insert(name, symbol);
    }

private:
    std::vector<std::shared_ptr<SymbolTable>> stack;
};

#endif
